from typing import*
import grpc

import knowledge_extract.service.pb.model_server_pb2 as pb
from knowledge_extract.service.pb.model_server_pb2_grpc import ModelServerStub
from knowledge_extract.service.tools import save_chunks_to_file
from knowledge_extract.service.utils.pb_helper import *


class ModelServerClient:
    def __init__(self, ip, port=50051):
        self.channel = grpc.insecure_channel(f'{ip}:{port}', options=[
            ('grpc.max_receive_message_length', 10 * 1024 * 1024),  # 设置为10MB
            ('grpc.max_send_message_length', 10 * 1024 * 1024),  # 设置为10MB
        ])
        self.stub = ModelServerStub(self.channel)

    def dataset_download(self, task, mode, dataset_name, output_file=None):
        request = pb.DatasetDownloadRequest(task=task, mode=mode, dataset_name=dataset_name)
        chunks = self.stub.dataset_download(request)
        if output_file is None:
            output_file = f'{task}_{mode}_{dataset_name}.json'
        save_chunks_to_file(chunks, output_file)
        assert os.path.exists(output_file)

    def create_experiment(self, task, model, dataset=None, experiment_id=None, project_id=None,
                          version=None,
                          model_config=None,
                          optimizer_config=None, data=None, label_list=None):
        experiment = PbHelper().create_experiment(task=task, model=model, dataset=dataset,
                                                  experiment_id=experiment_id,
                                                  project_id=project_id, version=version,
                                                  model_config=model_config,
                                                  optimizer_config=optimizer_config, data=data,
                                                  label_list=label_list)
        response = self.stub.create_experiment(experiment)
        return response

    def predict(self, text, task, model, dataset=None, version=None, experiment_id=None,
                project_id=None,
                model_config=None, label_list=None):
        request = PbHelper().create_predict_request(text=text, task=task, model=model,
                                                    dataset=dataset,
                                                    version=version,
                                                    experiment_id=experiment_id,
                                                    project_id=project_id,
                                                    model_config=model_config,
                                                    label_list=label_list)
        response = self.stub.predict(request)

        return response


    def run_active_learning(self, task, model, dataset, sampling_dataset, sampling_task, sampling_num=5, experiment_id=None, project_id=None,
                          version=None,
                          model_config=None,
                          optimizer_config=None, data=None, label_list=None):
        request = PbHelper().create_active_learning_request(task=task, model=model, dataset=dataset,
                                                  experiment_id=experiment_id,
                                                  project_id=project_id, version=version,
                                                  model_config=model_config,
                                                  optimizer_config=optimizer_config,
                                                  label_list=label_list,sampling_dataset=sampling_dataset, sampling_task=sampling_task, sampling_num=sampling_num)
        response = self.stub.run_active_learning(request)
        return response


    def hello(self, name):
        request = pb.HelloRequest(name=name)
        response = self.stub.hello(request)
        return response

    def close(self):
        self.channel.close()

    def get_running_metric(self, project_ids=None, experiment_ids=None):
        '''
        通过读取tfevent来获取指定实验的最新指标数据
        '''
        # 验证
        assert isinstance(project_ids, str) or project_ids is None
        assert isinstance(experiment_ids, str) or experiment_ids is None
        assert project_ids or experiment_ids

        # 设置缺省值
        project_ids = '-1' if not project_ids else project_ids
        experiment_ids = '-1' if not experiment_ids else experiment_ids

        # 构造请求
        request = pb.GetRunningMetricRequest(project_ids=project_ids, experiment_ids=experiment_ids)

        # 获取响应
        response = self.stub.get_running_metric(request)
        return response

    def save_dataset(self, dataset_name,train_dataset,dev_dataset,test_dataset,interface,task=None,file_extension=None):
        '''
        将数据集保存到数据集目录下
        '''
        # 验证
        assert task is None or isinstance(task, str)
        assert file_extension is None or isinstance(file_extension, str)
        assert isinstance(dataset_name, str)
        assert isinstance(train_dataset, bytes)
        assert isinstance(dev_dataset, bytes)
        assert isinstance(test_dataset, bytes)
        assert isinstance(interface, bytes)

        # 设置缺省值
        task = task or 'ner_re'
        file_extension = file_extension or '.json'

        # 构造请求
        request = pb.SaveRequest(task=task,dataset_name=dataset_name,train_dataset=train_dataset,dev_dataset=dev_dataset,test_dataset=test_dataset,interface=interface,dataset_extension=file_extension)

        # 获取响应
        response = self.stub.save_dataset(request)
        return response

    def get_model_params(self,model_name):
        '''
        通过读取config模板来获取参数项
        '''
        # 验证
        assert model_name

        # 构造请求
        request = pb.GetModelParamsRequest(model_name=model_name)

        # 获取响应
        response = self.stub.get_model_params(request)
        return response

    def recommend(self,task,dataset):
        assert isinstance(task,str)
        assert isinstance(dataset,str)

        request = pb.RecommendRequest(task=task,dataset=dataset)
        response = self.stub.recommend(request)
        return response